#!/usr/bin/env python3
# coding: utf-8

from numpy import array, ndarray, zeros, eye, round, deg2rad
from math import sin, cos

from robotics.io.pcd import load_point_cloud
from robotics.registration.util import transform_point_cloud, explain_quaternion
from robotics.registration.icp import Aligner
from robotics.registration.correspondence.oneway import find_correspondences
from robotics.registration.rotation.limity import extract_rotation
from robotics.registration.criteria.rms import stop_condition
from robotics.registration.merge.bruteforce import merge_point_cloud
from robotics.gui.glet import GridWindow
from robotics.state import StateMachine


def point_cloud_loader(prefix='/Users/yfwz100/ws/siasun/data/20160501/point'):
    for i in range(2, 10):
        p = load_point_cloud('%s%d.txt' % (prefix, i + 1))
        yield p / 10


window = GridWindow(caption='Eva')

point_clouds = point_cloud_loader()
source_cloud = target_cloud = None

aligner = Aligner(
        extract_rotation,
        find_correspondences,
    stop_condition(max_iterations=100, min_err=1, max_delta_err=1),
        precision=0
)
sm = StateMachine('source')
# window.time(500, sm.update)
window.key(sm.update)

# initial robot position and rotation.
first_pos = zeros((3,))
first_rot = eye(3)


def estimate_initial_position(r, t):
    global first_pos, first_rot
    first_pos = r @ (first_pos + t)
    first_rot = r @ first_rot
    print('-Init--------')
    print('Pos:\n%s' % round(first_pos, 2))
    print('Rot:\n%s,%s' % explain_quaternion(first_rot))
    print('-------------')


def rotation_around_y_estimation(deg):
    r = deg2rad(deg)
    return array([
        [cos(r), 0, -sin(r)],
        [0, 1, 0],
        [sin(r), 0, cos(r)]
    ])


@sm.state('source')
def pc_source(_):
    global source_cloud, counter
    source_cloud = next(point_clouds)
    return 'target'


@sm.state('target')
def pc_target(_):
    global target_cloud, counter
    try:
        target_cloud = next(point_clouds)
        return 'align'
    except FileNotFoundError and StopIteration:
        target_cloud = None
        return 'done'


@sm.state('align')
def pc_align(_):
    global source_cloud, target_cloud, first_pos, first_rot
    r, t = aligner.align(
            source_cloud.copy(),
            target_cloud.copy(),
        # rotation_around_y_estimation(3)
    )
    print("%.2f, %s" % explain_quaternion(r))
    source_cloud = transform_point_cloud(source_cloud, r, t)
    estimate_initial_position(r, t)
    return 'merge'


@sm.state('merge')
def pc_merge(_):
    global source_cloud, target_cloud
    source_cloud = merge_point_cloud(
            source_cloud,
            target_cloud,
            max_dist=10
    )
    print('Points: %d' % len(source_cloud))
    return 'target'


@window.within_draw_process
def draw(gl):
    global source_cloud, target_cloud

    gl.glPointSize(2)

    gl.glColor3d(0.1, 0.5, 1)

    if isinstance(source_cloud, ndarray):
        gl.glBegin(gl.GL_POINTS)
        for p in source_cloud:
            gl.glVertex3d(p[0] / 100, p[1] / 100, p[2] / 100)
        gl.glEnd()

    gl.glColor3d(0.1, 1, 0.5)

    if isinstance(target_cloud, ndarray):
        gl.glBegin(gl.GL_POINTS)
        for p in target_cloud:
            gl.glVertex3d(p[0] / 100, p[1] / 100, p[2] / 100)
        gl.glEnd()


window.run()
